#!/usr/bin/python
# -*- coding:utf-8 -*-

import subprocess
import time
import sys
import re
import argparse
import logging
import shlex

LOGFILE = "/var/log/clean_docker_images.log"
IMAGE_COMMAND_TEMPLATE = "docker image ls -a -q"
IMAGE_GREP_COMMAND_TEMPLATE = "docker image ls -a | grep %s"
IMAGE_CREATED_INFO_COMMAND_TEMPLATE = "docker inspect %s --format='{{ .Created }}'"
IMAGE_REMOVE_COMMAND_TEMPLATE = "docker rmi %s"


class Cmd(object):
    def __init__(self):
        self._return_code = None
        self._message = None

    def run(self, cmd_):
        if not isinstance(cmd_, str):
            msg = 'cmd is not available! expect Union[bytes, str] not %s' % type(cmd_)
            print(msg)
            sys.exit(1)
        cmd_list = shlex.split(cmd_)
        result = subprocess.Popen(cmd_list, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        result.wait()
        self._return_code = result.returncode
        self._message = result.stdout.read()
        # self._return_code, self._message = subprocess.getstatusoutput(cmd_)
        # if self._return_code != 0:
        #     msg = 'execute cmd error! error: %s' % self._message
        #     _print(msg)
        #     sys.exit(1)
        print('execute cmd: %s success' % cmd_)

    @property
    def return_code(self):
        return self._return_code

    @property
    def message(self):
        if isinstance(self._message, bytes):
            return self._message.decode(self, encoding='utf-8')
        return self._message


cmd = Cmd()


def command_rmi(remove_ids):
    for remove_id in remove_ids:
        # Replace shell pipe with Python filtering for SCA compliance
        cmd.run("docker image ls -a")
        if cmd.return_code == 0:
            # Filter lines containing remove_id in Python instead of using grep
            output_lines = cmd.message.split('\n')
            matching_lines = [line for line in output_lines if remove_id in line]
            if matching_lines:
                cmd.run(IMAGE_REMOVE_COMMAND_TEMPLATE % remove_id)
                if cmd.return_code == 0:
                    Logger(LOGFILE).get_log().info("%s" % cmd.message)


def transfer_timedelta(timedelta):
    if not re.match('^\d', timedelta):
        print("error format ! timedelta: %s " % timedelta)
        sys.exit(1)
    if timedelta.endswith('S'):
        return int(timedelta.split('S')[0])
    elif timedelta.endswith('M'):
        return int(timedelta.split('M')[0]) * 60
    elif timedelta.endswith('H'):
        return int(timedelta.split('H')[0]) * 60 * 60
    elif timedelta.endswith('d'):
        return int(timedelta.split('d')[0]) * 24 * 60 * 60
    elif timedelta.endswith('m'):
        return int(timedelta.split('m')[0]) * 30 * 24 * 60 * 60
    elif timedelta.endswith('y'):
        return int(timedelta.split('y')[0]) * 12 * 30 * 24 * 60 * 60
    else:
        print("error format ! timedelta: %s " % timedelta)
        sys.exit(1)


class Logger:
    def __init__(self, logger_name):
        self.logger = logging.getLogger()
        self.logger.setLevel(logging.INFO)  # Log等级总开关
        # 第二步，创建一个handler，用于写入日志文件
        logfile = logger_name
        fh = logging.FileHandler(logfile, mode='a')
        fh.setLevel(logging.INFO)  # 用于写到file的等级开关
        # 第三步，再创建一个handler,用于输出到控制台
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)  # 输出到console的log等级的开关
        # 第四步，定义handler的输出格式
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)
        ch.setFormatter(formatter)
        # 第五步，将logger添加到handler里面
        self.logger.addHandler(fh)
        self.logger.addHandler(ch)

    def get_log(self):
        return self.logger


class DockerImage(object):
    def __init__(self, ignore_ids=None):
        if ignore_ids is None:
            ignore_ids = []
        if not isinstance(ignore_ids, list):
            print("[ERROR] expect list!")
            sys.exit(1)
        self.ignore_ids = ignore_ids
        self.image_ids = []
        self.remove_ids = []
        self._get_image_ids()

    def _get_image_ids(self):
        cmd.run(IMAGE_COMMAND_TEMPLATE)
        if cmd.return_code != 0:
            print(
                '[ERROR] can not get docker image info! cmd: %s . please check docker status on your host' % IMAGE_COMMAND_TEMPLATE)
            sys.exit(1)
        self.image_ids = cmd.message.split('\n')[:-1]
        return self.image_ids

    def _check(self):
        pass


class DataParser(object):
    def parse(self, docker_image, timedelta):
        remove_ids = []
        image_id_string = ' '.join(docker_image.image_ids)
        timedelta = transfer_timedelta(timedelta)

        assert len(docker_image.image_ids) == len(self._get_created_timestamp_list(image_id_string))
        for image_id, created_timestamp in zip(docker_image.image_ids,
                                               self._get_created_timestamp_list(image_id_string)):
            if image_id in docker_image.ignore_ids:
                continue
            if created_timestamp <= time.time() - timedelta:
                remove_ids.append(image_id)
        docker_image.remove_ids = remove_ids

    def _get_created_timestamp_list(self, image_id_string):
        cmd.run(IMAGE_CREATED_INFO_COMMAND_TEMPLATE % image_id_string)
        if cmd.return_code != 0:
            msg = '[ERROR] can not get docker image id  info! cmd: %s . please check docker status on your host' % IMAGE_CREATED_INFO_COMMAND_TEMPLATE
            print(msg)
            sys.exit(1)
        return [self._transfer_timestring_to_timestamp(x) for x in cmd.message.split('\n')[:-1]]

    def _transfer_timestring_to_timestamp(self, timestring):
        return time.mktime(time.strptime(timestring.split(r'.')[0], "%Y-%m-%dT%H:%M:%S"))


data_parser = DataParser()
if __name__ == '__main__':

    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("-t", "--timedelta", default='3d', help="timedelta example 1S, 1M, 1H, 1d, 1m, 1y")
    arg_parser.add_argument("-g", "--ignore", nargs='*', help="ids in ignore_id_list will not be deleted")
    args = arg_parser.parse_args()

    docker_image = DockerImage(ignore_ids=args.ignore)
    data_parser.parse(docker_image, timedelta=args.timedelta)
    if len(docker_image.remove_ids) == 0:
        print("finish clean")
    command_rmi(docker_image.remove_ids)
